# 基于MNIST数据集，比较四种优化算法学习的快慢
import os
import sys
sys.path.append(os.pardir)  # 为了导入父目录的文件而进行的设定
import matplotlib.pyplot as plt
from dataset.mnist import load_mnist
from common.util import smooth_curve
from common.multi_layer_net import MultiLayerNet
from common.optimizer import *


# 0: 读入MNIST数据==================
(x_train,t_train),(x_test,t_test) = load_mnist(normalize=True)

train_size = x_train.shape[0]
batch_size = 128
max_iterations = 2000

# 1: 进行实验的设置===============
optimizers = {}
optimizers['SGD'] = SGD()
optimizers['Momentum'] = Momentum()
optimizers['AdaGrad'] = AdaGrad()
optimizers['Adam'] = Adam()

networks = {}
train_loss = {}
for key in optimizers.keys():
    networks[key] = MultiLayerNet(input_size=784,
                                  hidden_size_list=[100,100,100,100],
                                  output_size=10)
    train_loss[key]=[]

# 2:开始训练===========
for i in range(max_iterations):
    batch_mask = np.random.choice(train_size,batch_size)
    x_batch = x_train[batch_mask]
    t_batch = t_train[batch_mask]

    for key in optimizers.keys():
        grads = networks[key].gradient(x_batch,t_batch)
        optimizers[key].update(networks[key].params,grads)

        loss = networks[key].loss(x_batch,t_batch)
        train_loss[key].append(loss)

    if i % 100 == 0:
        print("===="+"iteration:"+str(i)+"====")
        for key in optimizers.keys():
            loss = networks[key].loss(x_batch,t_batch)
            print(key + " : "+str(loss))

# 3: 绘制图形==============
markers = {'SGD':"o","Momentum":"x","AdaGrad":"s","Adam":"D"}
x = np.arange(max_iterations)
for key in optimizers.keys():
    plt.plot(x,smooth_curve(train_loss[key]),marker=markers[key],markevery=100,label=key)
plt.xlabel("iterations")
plt.ylabel("loss")
plt.ylim(0,1)
plt.legend()
plt.show()